import os

import tqdm

from gxl_ai_utils.utils import utils_file


def do_split(input_wav_scp_path, input_text_scp_path, output_dir, split_num):
    """
    做切割，同时筛选出text为空的数据
    :param input_wav_scp_path:
    :param input_text_scp_path:
    :param output_dir:
    :param split_num:
    :return:
    """
    utils_file.logging_limit_print('start do_split')
    now = utils_file.do_get_now_time()
    wav_dict = utils_file.load_dict_from_scp(input_wav_scp_path)
    text_dict = utils_file.load_dict_from_scp(input_text_scp_path)

    utils_file.logging_print(f'wav_dict: {len(wav_dict)}')
    utils_file.logging_print(f'text_dict: {len(text_dict)}')
    wav_dict_list = utils_file.do_split_dict(wav_dict, split_num)
    for i in tqdm.tqdm(range(split_num), desc='do_split for scp', total=split_num):
        output_dir_i = os.path.join(output_dir, f'temp_{i}')
        utils_file.makedir_sil(output_dir_i)
        wav_dict_i = wav_dict_list[i]
        text_dict_i = {k: v for k, v in text_dict.items() if k in wav_dict_i}
        # 筛选text
        text_dict_i_filter = {k: v for k, v in text_dict_i.items() if len(v.strip()) > 0}
        wav_dict_i_filter = {k: v for k, v in wav_dict_i.items() if k in text_dict_i_filter}
        output_wav_scp_path_i = os.path.join(output_dir_i, 'wav.scp')
        output_text_scp_path_i = os.path.join(output_dir_i, 'text')
        utils_file.write_dict_to_scp(wav_dict_i_filter, output_wav_scp_path_i)
        utils_file.write_dict_to_scp(text_dict_i_filter, output_text_scp_path_i)
    utils_file.logging_limit_print(f'finish do_split, 耗时:{utils_file.do_get_elapsed_time(now)}s')

def do_concat_some_scp():
    """"""
    row_pachong_scp_dir = '/home/work_nfs8/xlgeng/aslp_spider_data'
    utils_file.makedir_sil(row_pachong_scp_dir)
    input_dirs = [
        # "/home/work_nfs8/yhdai/data01/no_vad_processed/xmly_entertainment_1",
        # "/home/work_nfs8/yhdai/data01/no_vad_processed/yunting_jingjizhisheng",
        # "/home/work_nfs8/yhdai/data01/no_vad_processed/xmly10T_jiankang_2527",
        "/home/work_nfs8/xlgeng/aslp_spider_data/npr_en"
    ]
    big_name= 'pici_26775H'
    secend_name = 'main_8_6300H'
    output_dir_now = os.path.join(row_pachong_scp_dir, big_name, secend_name)
    utils_file.makedir_sil(output_dir_now)
    res_wav_dict = {}
    res_txt_dict = {}

    for input_dir in input_dirs:
        input_wav_scp_path = os.path.join(input_dir, 'wav.scp')
        input_text_scp_path = os.path.join(input_dir, 'text')
        if not os.path.exists(input_wav_scp_path) or not os.path.exists(input_text_scp_path):
            utils_file.logging_print(f'input_wav_scp_path= {input_wav_scp_path} or input_text_scp_path= {input_text_scp_path} not exists')
            continue
        wav_dict = utils_file.load_dict_from_scp(input_wav_scp_path, silence=True)
        txt_dict = utils_file.load_dict_from_scp(input_text_scp_path, silence=True)
        res_wav_dict.update(wav_dict)
        res_txt_dict.update(txt_dict)
    # do filter
    res_txt_dict_filter = {k: v for k, v in res_txt_dict.items() if len(v.strip()) > 0}
    res_wav_dict_filter = {k: v for k, v in res_wav_dict.items() if k in res_txt_dict_filter}
    output_wav_scp_path = os.path.join(output_dir_now, 'wav.scp')
    output_text_scp_path = os.path.join(output_dir_now, 'text')
    utils_file.write_dict_to_scp(res_wav_dict_filter, output_wav_scp_path)
    utils_file.write_dict_to_scp(res_txt_dict_filter, output_text_scp_path)
    utils_file.logging_print(f'finish , output_dir_now= {output_dir_now}')
    nodes = [
        'node1',
    ]
    res_wav_dict_filter_list = utils_file.do_split_dict(res_wav_dict_filter,len(nodes))

    for i, node in enumerate(nodes):
        res_wav_dict_filter_i = res_wav_dict_filter_list[i]
        res_txt_dict_filter_i = {k: res_txt_dict_filter[k] for k in res_wav_dict_filter_i.keys()}
        output_dir_now_i = os.path.join(output_dir_now, node)
        utils_file.makedir_sil(output_dir_now_i)
        output_wav_scp_path_i = os.path.join(output_dir_now_i, 'wav.scp')
        output_text_scp_path_i = os.path.join(output_dir_now_i, 'text')
        utils_file.write_dict_to_scp(res_wav_dict_filter_i, output_wav_scp_path_i)
        utils_file.write_dict_to_scp(res_txt_dict_filter_i, output_text_scp_path_i)
        utils_file.logging_print(f'finish , output_dir_now_i= {output_dir_now_i}')


def little_func_for_tts_to_scp(tts_dict,index_i, output_dir):
    output_dir_now = os.path.join(output_dir, 'split', f'split_{index_i}')
    utils_file.makedir_sil(output_dir_now)
    wav_dict = {}
    text_dict = {}
    for key, value in tqdm.tqdm(tts_dict.items()):
        wav_path, text_path = value.strip().split()
        wav_dict[key] = wav_path
        text_info = utils_file.load_first_row_clean(text_path)
        res_text = text_info.strip().split()
        text_key = res_text[0]
        text = ' '.join(res_text[1:])
        # print(text_key, text)
        text_dict[text_key] = text
    output_wav_scp_path_i = os.path.join(output_dir_now, 'wav.scp')
    output_text_scp_path_i = os.path.join(output_dir_now, 'text')
    utils_file.write_dict_to_scp(wav_dict, output_wav_scp_path_i)
    utils_file.write_dict_to_scp(text_dict, output_text_scp_path_i)
def convert_tts_to_scp():
    input_tts_path = "/home/work_nfs10/kxxia/processed_data/liukai_NPR-en_final.lst"
    tts_dict = utils_file.load_dict_from_scp(input_tts_path)
    tts_didct_list = utils_file.do_split_dict(tts_dict, 20)
    runner = utils_file.GxlDynamicThreadPool()
    output_wav_scp_dir = "/home/work_nfs8/xlgeng/aslp_spider_data/npr_en"
    index = 0
    for tts_dict_i in tts_didct_list:
        runner.add_thread(little_func_for_tts_to_scp, [tts_dict_i, index, output_wav_scp_dir])
        index += 1
    runner.run()
    utils_file.makedir_sil(output_wav_scp_dir)
    output_wav_scp_path = os.path.join(output_wav_scp_dir, 'wav.scp')
    output_text_scp_path = os.path.join(output_wav_scp_dir, 'text')
    res_wav_dict = {}
    res_txt_dict = {}
    for i in range(index):
        output_dir_now = os.path.join(output_wav_scp_dir, 'split', f'split_{i}')
        temp_wav_path = utils_file.load_dict_from_scp(os.path.join(output_dir_now, 'wav.scp'))
        temp_text_path = utils_file.load_dict_from_scp(os.path.join(output_dir_now, 'text'))
        temp_wav_dict = utils_file.load_dict_from_scp(temp_wav_path)
        temp_text_dict = utils_file.load_dict_from_scp(temp_text_path)
        res_wav_dict.update(temp_wav_dict)
        res_txt_dict.update(temp_text_dict)

    utils_file.write_dict_to_scp(res_wav_dict, output_wav_scp_path)
    utils_file.write_dict_to_scp(res_txt_dict, output_text_scp_path)
    res = utils_file.do_execute_shell_command(f"head -n 100 {output_wav_scp_path}")
    print(res)
    res = utils_file.do_execute_shell_command(f"head -n 100 {output_text_scp_path}")
    print(res)

if __name__ == '__main__':
    convert_tts_to_scp()
    do_concat_some_scp()